package common

type TaskPool struct {
	numWorks int
	taskChan chan Task
	stopChan chan struct{}
}

type Task struct {
	Fu func()
}

func NewTaskPool(numWorks int) *TaskPool {
	return &TaskPool{
		numWorks: numWorks,
		taskChan: make(chan Task),
		stopChan: make(chan struct{}),
	}
}

func (p *TaskPool) worker() {
	for {
		select {
		case task, ok := <-p.taskChan:
			if !ok {
				return
			}
			task.Fu()
		case <-p.stopChan:
			return
		}
	}
}

func (p *TaskPool) Start() {
	for i := 0; i < p.numWorks; i++ {
		go p.worker()
	}
}

func (p *TaskPool) Stop() {
	close(p.taskChan)
	close(p.stopChan)
}

func (p *TaskPool) AddTask(task Task) {
	p.taskChan <- task
}
